/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.gluten.extension

import org.apache.gluten.config.GlutenConfig
import org.apache.gluten.expression.aggregate.HLLAdapter

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.Literal
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, HyperLogLogPlusPlus}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, AGGREGATE_EXPRESSION}
import org.apache.spark.sql.types._

case class HLLRewriteRule(spark: SparkSession) extends Rule[LogicalPlan] {
  override def apply(plan: LogicalPlan): LogicalPlan = {
    plan.transformUpWithPruning(_.containsPattern(AGGREGATE)) {
      case a: Aggregate =>
        a.transformExpressionsWithPruning(_.containsPattern(AGGREGATE_EXPRESSION)) {
          case aggExpr @ AggregateExpression(hll: HyperLogLogPlusPlus, _, _, _, _)
              if GlutenConfig.get.enableNativeHyperLogLogAggregateFunction &&
                GlutenConfig.get.enableColumnarHashAgg &&
                isSupportedDataType(hll.child.dataType) =>
            val hllAdapter = HLLAdapter(
              hll.child,
              Literal(hll.relativeSD),
              hll.mutableAggBufferOffset,
              hll.inputAggBufferOffset)

            aggExpr.copy(aggregateFunction = hllAdapter)
        }
    }
  }

  private def isSupportedDataType(dataType: DataType): Boolean = {
    // HLL in Velox only supports below data types. We should not offload HLL to velox, if
    // child's data type is not supported. This prevents the case only partail agg is fallbacked.
    // As Spark and Velox have different HLL binary formats, HLL binary generated by Spark can't
    // be parsed by Velox, it would cause the error: 'Unexpected type of HLL'.
    dataType match {
      case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
          StringType | _: CharType | _: DecimalType | DateType | TimestampType | BinaryType =>
        true
      case _ => false
    }
  }
}
